#include <limits>
#include <cstdlib>
#include <ctime>
#include "readfile.cpp"

struct Centroid {
    double latitude;
    double longitude;
    double averageSignal;
    double range;
};

double distance(const Point& p, const Centroid& c) {
    return std::sqrt(
        std::pow(p.latitude - c.latitude, 2) +
        std::pow(p.longitude - c.longitude, 2) +
        std::pow(p.averageSignal - c.averageSignal, 2) +
        std::pow(p.range - c.range, 2)
    );
}

void kmeans(std::vector<Point>& data, int k, int maxIterations, bool parallel) {
    int n = data.size();
    std::vector<Centroid> centroids(k);
    std::srand(std::time(nullptr));

    // Initialize centroids randomly
    for (int i = 0; i < k; ++i) {
        int index = std::rand() % n;
        centroids[i].latitude = data[index].latitude;
        centroids[i].longitude = data[index].longitude;
        centroids[i].averageSignal = data[index].averageSignal;
        centroids[i].range = data[index].range;
    }

    for (int iter = 0; iter < maxIterations; ++iter) {
        // Assign points to the nearest centroid
        if (parallel) {
            #pragma omp parallel for
            for (int i = 0; i < n; ++i) {
                double minDist = std::numeric_limits<double>::max();
                int bestCluster = 0;
                for (int j = 0; j < k; ++j) {
                    double dist = distance(data[i], centroids[j]);
                    if (dist < minDist) {
                        minDist = dist;
                        bestCluster = j;
                    }
                }
                data[i].cluster = bestCluster;
            }
        } else {
            for (int i = 0; i < n; ++i) {
                double minDist = std::numeric_limits<double>::max();
                int bestCluster = 0;
                for (int j = 0; j < k; ++j) {
                    double dist = distance(data[i], centroids[j]);
                    if (dist < minDist) {
                        minDist = dist;
                        bestCluster = j;
                    }
                }
                data[i].cluster = bestCluster;
            }
        }

        // Temporary arrays for parallel reduction
        std::vector<Centroid> newCentroids(k, {0, 0, 0, 0});
        std::vector<int> counts(k, 0);

        if (parallel) {
            #pragma omp parallel
            {
                std::vector<Centroid> localCentroids(k, {0, 0, 0, 0});
                std::vector<int> localCounts(k, 0);

                #pragma omp for nowait
                for (int i = 0; i < n; ++i) {
                    int cluster = data[i].cluster;
                    localCentroids[cluster].latitude += data[i].latitude;
                    localCentroids[cluster].longitude += data[i].longitude;
                    localCentroids[cluster].averageSignal += data[i].averageSignal;
                    localCentroids[cluster].range += data[i].range;
                    localCounts[cluster] += 1;
                }

                #pragma omp critical
                {
                    for (int j = 0; j < k; ++j) {
                        newCentroids[j].latitude += localCentroids[j].latitude;
                        newCentroids[j].longitude += localCentroids[j].longitude;
                        newCentroids[j].averageSignal += localCentroids[j].averageSignal;
                        newCentroids[j].range += localCentroids[j].range;
                        counts[j] += localCounts[j];
                    }
                }
            }
        } else {
            for (int i = 0; i < n; ++i) {
                int cluster = data[i].cluster;
                newCentroids[cluster].latitude += data[i].latitude;
                newCentroids[cluster].longitude += data[i].longitude;
                newCentroids[cluster].averageSignal += data[i].averageSignal;
                newCentroids[cluster].range += data[i].range;
                counts[cluster] += 1;
            }
        }

        for (int j = 0; j < k; ++j) {
            if (counts[j] > 0) {
                centroids[j].latitude = newCentroids[j].latitude / counts[j];
                centroids[j].longitude = newCentroids[j].longitude / counts[j];
                centroids[j].averageSignal = newCentroids[j].averageSignal / counts[j];
                centroids[j].range = newCentroids[j].range / counts[j];
            }
        }
    }
}
